# 给你一个 n x n 矩阵 matrix ，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。
# 请注意，它是 排序后 的第 k 小元素，而不是第 k 个 不同 的元素。
#
#  示例 1：
# 输入：matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
# 输出：13
# 解释：矩阵中的元素为 [1,5,9,10,11,12,13,13,15]，第 8 小元素是 13
#
#  示例 2：
# 输入：matrix = [[-5]], k = 1
# 输出：-5
import heapq
from typing import List


class Solution:
    def kthSmallest2(self, matrix: List[List[int]], k: int) -> int:
        """
        二分查找
        矩阵内的元素是从左上到右下递增的，matrix[0][0], matrix[-1][-1]分别为最小值、最大值记为left,right
        任取一个数 mid 满足 left <= mid <= right，那么矩阵中不大于mid 的数，肯定全部分布在矩阵的左上角。

        矩阵中大于 mid 的数就和不大于 mid 的数分别形成了两个板块，沿着一条锯齿线将这个矩形分开。其中左上角板块的大小即为矩阵中不大于 mid 的数的数量。
        当左上角板块的大小等于了k - 1就找到了 矩阵中的第K小的元素
        [[1, 5, 9],
        [10, 11, 13],
        [12, 13, 15]]

        :param matrix:
        :param k:
        :return:
        """
        rows, columns = len(matrix), len(matrix[0])
        left, right = matrix[0][0], matrix[-1][-1]

        def check(target: int) -> bool:  # 检查矩阵中小于target的数的数量是否大于等于k
            i, j = rows - 1, 0
            count = 0
            while i >= 0 and j < columns:  # 从左下角开始统计比target小的数的数量
                if matrix[i][j] <= target:
                    count += i + 1
                    j += 1
                else:
                    i -= 1
            return count >= k

        while left < right:
            mid = left + (right - left) // 2
            if check(mid):
                right = mid
            else:
                left = mid + 1
        return left

    def kthSmallest1(self, matrix: List[List[int]], k: int) -> int:
        """
        归并排序合并的过程
        由于有多个有序的数组(每一行可以看成数组)，所以可以用最小堆来优化
        合并到k - 1次
        :param matrix:
        :param k:
        :return:
        """
        heap = [(matrix[i][0], i, 0) for i in range(len(matrix))]
        heapq.heapify(heap)
        for _ in range(k - 1):
            num, row, column = heapq.heappop(heap)
            if column < len(matrix[0]) - 1:
                heapq.heappush(heap, (matrix[row][column + 1], row, column + 1))
        return heapq.heappop(heap)[0]

    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        return self.kthSmallest2(matrix, k)


if __name__ == "__main__":
    matrix = [[1, 5, 9],
              [10, 11, 13],
              [12, 13, 15]]
    k = 8
    print(Solution().kthSmallest(matrix, k))
